# coding: utf8
import os
import random
import re
from typing import List

import dotenv
from openai import OpenAI

# print(os.getenv("DeepSeek"))
dotenv.load_dotenv(".env")
# print(os.getenv("DeepSeek"))
from logger import Logger


class LLMBase:

    # SYS_PROMPT = """你是一个土木工程领域的专家，需要为构建土木抗震领域大语言模型的微调数据集提供帮助，用户会给你提供从抗震相关教材或规范中截取的内容(其中公式为latex)，为避免单次内容超过长度限制，某章内容将分为多次对话提供。用户提供材料后，要求你：1.请从中总结出知识点，形成问答对；2.问答对以纯文本形式返回，格式如下：Q:问题1\nA:答案1\n\nQ:问题2\nA:答案2...；3.若问题或回答涉及公式，以latex格式表示。由于材料中图片及表格被去除，问题和答案中不应提及某图或某表；4.随着上下文语境增长，构建问答时应考虑多轮对话所给材料之间知识的联系，而不仅仅从当前对话所给材料来总结，以保证知识问答对的全面性，且使的解答足够详尽清晰；5.对于问题，内容中不要包含如“思考题，习题”之类的前置词，仅包含纯问题；对于答案，需要尽量详细具体（重要！）,提供清晰解答后，需结合相关知识进一步展开以补充细节或解释；6.尽可能多地总结出Q&A，越多越好（重要！），使得它们覆盖给定的材料，可以尝试从不同的角度设问。若内容无法总结出有效问答，则返回空列表[]；7.在详细回答问题的同时需要保证其准确性和专业性，不能杜撰。8. 注意返回的内容只包含Q&A对，切勿包含其他内容。"""
    SYS_PROMPT = """你是一个土木工程抗震研究领域的专家"""
    PROMPT = """我需要构建语言大模型的微调数据集，这个大模型面向的领域是土木工程专业，因此我搜集了大量书籍，下面我会给你提供从中截取的内容(其中公式为latex)，要求：1.请从中总结出有效的知识点；2.请以问答对的纯文本形式返回，格式如下：Q:问题1\nA:答案1\n\nQ:问题2\nA:答案2...；3.若问题或回答涉及公式，以latex格式表示。问题或回答中不能引用某图或某表，因为本模型无法生成图片或表格；4.答案需要尽量详细具体（重要！）、逻辑清晰、回答全面，其中应避免提及原材料中可能存在的图片，如不要出现“如图xxx所示”类似话语，即不要通过引用图片来进行提问或解释；5.尽可能多地总结出Q&A，越多越好（重要！），使得它们覆盖给定的材料，可以尝试从不同的角度设问。若内容无法总结出有效问答，则返回空列表[]；6.在详细回答问题的同时需要保证其准确性和专业性，不能杜撰。7. 注意返回的内容只包含Q&A对，切勿包含其他内容。\n"""

    # 效果不错
    # PROMPT = """我需要构建语言大模型的微调数据集，这个大模型面向的领域是土木工程专业，因此我搜集了大量书籍，下面我会给你提供从中截取的内容(其中公式为latex)，要求：1.请从中总结出有效的知识点；2.请以问答对的纯文本形式返回，格式如下：Q:问题1\nA:答案1\n\nQ:问题2\nA:答案2...；3.若问题或回答涉及公式，以latex格式表示；4.答案需要尽量详细丰富（重要！）；5.尽力总结出更多问答对，越多越好（重要！），使得它们覆盖给定的材料，可以尝试从不同的角度设问。若内容无法总结出有效问答，则返回空列表[]；6.在详细回答问题的同时需要保证其准确性和专业性，不能杜撰。7. 注意返回的内容只包含Q&A对，切勿包含其他内容。
    # 以下是材料："""

    def __init__(self, llm_name: str, logger: Logger = None, api_key=None) -> None:
        self.llm_name = llm_name
        self.logger = logger if logger else Logger("chat.log")
        self.history = []  # 用于存储历史对话
        self.curent_chapter: str = None

        self.default_kwargs = {}
        self.client = self.create_client(api_key)

    def create_client(self, api_key=None) -> OpenAI:
        if api_key is None:
            api_key = os.getenv(self.llm_name)
            # print("API_KEY:", api_key)

        if not api_key:
            raise ValueError("API_KEY is not set for " + self.llm_name)

        # print("API_KEY:", api_key)
        self._api_key = api_key
        if self.llm_name == "ChatGPT":
            client = OpenAI(api_key=api_key)
            self.default_kwargs["model"] = "gpt-3.5-turbo"
        elif self.llm_name == "DeepSeek":
            client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
            self.default_kwargs["model"] = "deepseek-chat"
        elif self.llm_name == "ZeroOne":
            client = OpenAI(api_key=api_key, base_url="https://api.lingyiwanwu.com/v1")
            self.default_kwargs["model"] = "yi-large"
        else:
            raise ValueError(
                "The LLM is not supported yet:" + self.llm_name,
                "available LLMs are: [ChatGPT, DeepSeek, ZeroOne]",
            )

        return client

    def parse_qa_result(self, response):
        """
        parse Q&A pairs from LLM response and return a list of Q&A pairs
        """
        text_res = response.choices[0].message.content
        with open("response.txt", "a", encoding="utf-8") as f:
            f.write(text_res + "\n")

        questions: List[str] = re.findall(r"(?<=Q:).+?(?=\nA|$)", text_res, re.DOTALL)
        answers: List[str] = re.findall(r"(?<=A:).+?(?=\nQ|$)", text_res, re.DOTALL)
        if len(questions) != len(answers):
            self.logger.log(
                "Questions and Answers are not in pairs.\n"
                + "questions: "
                + str(questions)
                + "\n"
                + "answers: "
                + str(answers)
            )
            return []

        # create a list of Q&A pairs
        qa_pairs = []
        for q, a in zip(questions, answers):
            qa_pairs.append({"Q": q.strip(), "A": a.strip()})
        return qa_pairs

    def chat(self, messages, sys_prompt=None, **kwargs):
        """
        send messages to LLM and get response
        Args:
         - messages: a list of messages history in format [{"role": "user|assistant", "content": "text"}]
           or a single string message to be sent to ChatGPT
        return raw ChatCompletion object
        """
        response = None
        kwargs = {**self.default_kwargs, **kwargs}
        if sys_prompt is None:
            sys_prompt = self.SYS_PROMPT

        if isinstance(messages, str):
            messages = [{"role": "user", "content": messages}]

        response = self.client.chat.completions.create(
            messages=[
                {"role": "system", "content": sys_prompt},
                *messages,
            ],
            **kwargs,
        )
        return response

    def gen_qa_single_round(self, text, **kwargs):
        """
        Args:
         - text: the text to be sent to ChatGPT
         - use_prefix: whether to use the default prefix, act as a q&a pair generator
        """
        messages = [{"role": "user", "content": self.PROMPT + text}]
        response = self.chat(messages, **kwargs)
        if not response:
            return []

        qa_pairs = self.parse_qa_result(response)
        return qa_pairs

    def gen_qa_multi_round(self, text, max_context_len, start_new=False, **kwargs):
        """
        for multi-round chat

        Args:
         - text: the text to be sent to ChatGPT
         - max_context_len: the context length limit for the model
         - start_new: whether to start a new chat
        """

        # clear history if start a new chat
        if start_new:
            self.history = []

        messages = [*self.history, {"role": "user", "content": text}]
        # 发送请求
        response = self.chat(messages, **kwargs)
        # 将当前对话添加到历史记录中
        self.history.append({"role": "user", "content": text})
        self.history.append({"role": "assistant", "content": response.choices[0].message.content})

        qa_pairs = self.parse_qa_result(response)

        # 计算当前历史记录的总长度
        total_length = response.usage.total_tokens
        # 如果总长度超过限值，删除最早的消息及回复
        if total_length > max_context_len:
            self.history.pop(0)
            self.history.pop(0)

        return qa_pairs

    def get_error_code(self, error):
        """
        get error code from the error message like:Error code: 400
        """
        err_code = re.search(r"(?<=Error code: )\d{3}", str(error))
        if err_code:
            return err_code.group()
        else:
            return "Unknown Error"

    def record_response(self, response):
        role = response.choices[0].message.role
        content = response.choices[0].message.content
        self.logger.log(f"{role}: {content}\n\n\n")
        return content

    def change_key(self):
        """
        keep a api key pool for DeepSeek
        """
        if self.llm_name != "DeepSeek":
            return
        with open("deepseek_keys.txt", "r", encoding="utf-8") as f:
            api_keys = [line.strip("\n") for line in f.readlines()]

        if not hasattr(self, "usedup_keys"):
            self.usedup_keys = []

        self.usedup_keys.append(self._api_key)
        with open("usedup_keys.txt", "a", encoding="utf-8") as f:
            f.write(self._api_key + "\n")

        api_keys = list(set(api_keys) - set(self.usedup_keys))
        if not api_keys:
            self.logger.log("All keys are used up.")
            return

        api_key = random.choice(api_keys)
        self.client = self.create_client(api_key)
